package com.zeyu.framework.core.configuration;

import com.zeyu.framework.tools.console.HTTPTunnelRequest;
import com.zeyu.framework.tools.console.TunnelRequest;
import com.zeyu.framework.tools.console.TunnelRequestService;
import com.zeyu.framework.utils.SpringContextHolder;
import org.apache.guacamole.GuacamoleClientException;
import org.apache.guacamole.GuacamoleConnectionClosedException;
import org.apache.guacamole.GuacamoleException;
import org.apache.guacamole.io.GuacamoleReader;
import org.apache.guacamole.io.GuacamoleWriter;
import org.apache.guacamole.net.GuacamoleTunnel;
import org.apache.guacamole.protocol.GuacamoleInstruction;
import org.apache.guacamole.protocol.GuacamoleStatus;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.context.annotation.Configuration;
import org.springframework.messaging.simp.config.MessageBrokerRegistry;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.config.annotation.AbstractWebSocketMessageBrokerConfigurer;
import org.springframework.web.socket.config.annotation.EnableWebSocket;
import org.springframework.web.socket.config.annotation.EnableWebSocketMessageBroker;
import org.springframework.web.socket.config.annotation.StompEndpointRegistry;
import org.springframework.web.socket.config.annotation.WebSocketConfigurer;
import org.springframework.web.socket.config.annotation.WebSocketHandlerRegistry;
import org.springframework.web.socket.handler.TextWebSocketHandler;
import org.springframework.web.socket.server.support.DefaultHandshakeHandler;

import java.io.IOException;

/**
 * 配置websocket的支持,为了浏览器兼容性,我们使用sockJS
 * 目前没有统计需求,直接使用这个就可以了
 * Created by zeyuphoenix on 16/10/11.
 */
@Configuration
@EnableWebSocketMessageBroker
@EnableWebSocket
public class WebSocketConfiguration extends AbstractWebSocketMessageBrokerConfigurer implements WebSocketConfigurer {

    // ================================================================
    // Constants
    // ================================================================

    /**
     * logger
     */
    private static final Logger logger = LoggerFactory.getLogger(WebSocketConfiguration.class);

    // ================================================================
    // Fields
    // ================================================================

    // ================================================================
    // Constructors
    // ================================================================

    // ================================================================
    // Methods from/for super Interfaces or SuperClass
    // ================================================================

    /**
     * 添加一个服务端点，来接收客户端的连接
     */
    @Override
    public void registerStompEndpoints(StompEndpointRegistry registry) {
        // 表示添加了一个/socket端点，客户端就可以通过这个端点来进行连接
        registry.addEndpoint("/socket")
                // 作用是开启SockJS支持
                .withSockJS()
                // 设置为512K
                .setStreamBytesLimit(512 * 1024)
                .setHttpMessageCacheSize(1000)
                .setDisconnectDelay(30 * 1000L);
    }

    /**
     * 定义消息代理，通俗一点讲就是设置消息连接请求的各种规范信息
     */
    @Override
    public void configureMessageBroker(MessageBrokerRegistry registry) {
        // 表示客户端订阅地址的前缀信息，也就是客户端接收服务端消息的地址的前缀信息
        registry.enableSimpleBroker("/topic");
        // 服务端接收地址的前缀，意思就是说客户端给服务端发消息的地址的前缀
        registry.setApplicationDestinationPrefixes("/logger");
    }

    /**
     * 支持普通的websocket,目前主要用于console
     */
    @Override
    public void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {

        // set Handshake protocol
        DefaultHandshakeHandler handshakeHandler = new DefaultHandshakeHandler();
        handshakeHandler.setSupportedProtocols("guacamole");

        // add handler
        registry.addHandler(new ConsoleTextWebSocketHandler(), "/console/websocket-tunnel")
                .setHandshakeHandler(handshakeHandler);
    }

    // ================================================================
    // Public or Protected Methods
    // ================================================================

    // ================================================================
    // Getter & Setter
    // ================================================================

    // ================================================================
    // Private Methods
    // ================================================================

    // ================================================================
    // Inner or Anonymous Class
    // ================================================================

    private class ConsoleTextWebSocketHandler extends TextWebSocketHandler {

        /**
         * The default, minimum buffer size for instructions.
         */
        private static final int BUFFER_SIZE = 8192;

        /**
         * The underlying GuacamoleTunnel. WebSocket reads/writes will be handled
         * as reads/writes to this tunnel.
         */
        private GuacamoleTunnel tunnel;

        @Override
        protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
            super.handleTextMessage(session, message);
            // Ignore inbound messages if there is no associated tunnel
            if (tunnel == null)
                return;

            GuacamoleWriter writer = tunnel.acquireWriter();

            try {
                // Write received message
                if (message != null) {
                    String payload = message.getPayload();
                    if (payload != null) {
                        writer.write(payload.toCharArray());
                    }
                }
            } catch (GuacamoleConnectionClosedException e) {
                logger.debug("Connection to guacd closed.", e);
            } catch (GuacamoleException e) {
                logger.debug("WebSocket tunnel write failed.", e);
            }

            tunnel.releaseWriter();
        }

        @Override
        public void afterConnectionEstablished(WebSocketSession session) throws Exception {
            super.afterConnectionEstablished(session);
            logger.debug("console session is connection ");

            try {

                // get tunnel request
                TunnelRequest request = new HTTPTunnelRequest(session.getUri().getQuery());
                // Get tunnel, 因为创建顺序,不能注入
                TunnelRequestService tunnelRequestService = SpringContextHolder.getBean(TunnelRequestService.class);
                tunnel = tunnelRequestService.createTunnel(request);
                if (tunnel == null) {
                    closeConnection(session, GuacamoleStatus.RESOURCE_NOT_FOUND);
                    return;
                }

            } catch (GuacamoleException e) {
                logger.error("Creation of WebSocket tunnel to guacd failed: {}", e.getMessage());
                logger.debug("Error connecting WebSocket tunnel.", e);
                closeConnection(session, e.getStatus());
                return;
            }

            // Prepare read transfer thread
            Thread readThread = new Thread() {

                /**
                 * Remote (client) side of this connection
                 */
                @Override
                public void run() {

                    StringBuilder buffer = new StringBuilder(BUFFER_SIZE);
                    GuacamoleReader reader = tunnel.acquireReader();
                    char[] readMessage;

                    try {

                        // Send tunnel UUID
                        session.sendMessage(new TextMessage(new GuacamoleInstruction(
                                GuacamoleTunnel.INTERNAL_DATA_OPCODE,
                                tunnel.getUUID().toString()
                        ).toString()));

                        try {

                            // Attempt to read
                            while ((readMessage = reader.read()) != null) {

                                // Buffer message
                                buffer.append(readMessage);

                                // Flush if we expect to wait or buffer is getting full
                                if (!reader.available() || buffer.length() >= BUFFER_SIZE) {
                                    session.sendMessage(new TextMessage(buffer.toString()));
                                    buffer.setLength(0);
                                }

                            }

                            // No more data
                            closeConnection(session, GuacamoleStatus.SUCCESS);

                        }
                        // Catch any thrown guacamole exception and attempt
                        // to pass within the WebSocket connection, logging
                        // each error appropriately.
                        catch (GuacamoleClientException e) {
                            logger.info("WebSocket connection terminated: {}", e.getMessage());
                            logger.debug("WebSocket connection terminated due to client error.", e);
                            closeConnection(session, e.getStatus());
                        } catch (GuacamoleConnectionClosedException e) {
                            logger.debug("Connection to guacd closed.", e);
                            closeConnection(session, GuacamoleStatus.SUCCESS);
                        } catch (GuacamoleException e) {
                            logger.error("Connection to guacd terminated abnormally: {}", e.getMessage());
                            logger.debug("Internal error during connection to guacd.", e);
                            closeConnection(session, e.getStatus());
                        }

                    } catch (IOException e) {
                        logger.debug("I/O error prevents further reads.", e);
                        closeConnection(session, GuacamoleStatus.SERVER_ERROR);
                    }

                }

            };

            readThread.start();
        }

        @Override
        public void handleTransportError(WebSocketSession session, Throwable exception) throws Exception {
            super.handleTransportError(session, exception);
            logger.debug("WebSocket tunnel closing due to error", exception);
            try {
                if (tunnel != null)
                    tunnel.close();
            } catch (GuacamoleException e) {
                logger.debug("Unable to close connection to guacd.", e);
            }
        }

        @Override
        public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
            super.afterConnectionClosed(session, status);
            logger.debug("close console websocket tunnel, status is {}", status);
            try {
                if (tunnel != null)
                    tunnel.close();
            } catch (GuacamoleException e) {
                logger.debug("Unable to close connection to guacd.", e);
            }

        }

        /**
         * Sends the given status on the given WebSocket connection and closes the
         * connection.
         *
         * @param session The outbound WebSocket connection to close.
         * @param guacStatus The status to send.
         */
        private void closeConnection(WebSocketSession session, GuacamoleStatus guacStatus) {

            try {
                int code = guacStatus.getWebSocketCode();
                String message = Integer.toString(guacStatus.getGuacamoleStatusCode());
                session.close(new CloseStatus(code, message));
            } catch (IOException e) {
                logger.debug("Unable to close WebSocket connection.", e);
            }
        }
    }

    // ================================================================
    // Test Methods
    // ================================================================

}
